import os
import requests
import m3u8
import concurrent.futures
import subprocess
from tqdm import tqdm
from urllib.parse import urljoin

def download_ts_segment(url, output_path, max_retries=3):
    """
    Downloads a single .ts segment with retry mechanism.
    """
    for attempt in range(max_retries):
        try:
            response = requests.get(url, stream=True, timeout=10)
            response.raise_for_status()
            with open(output_path, 'wb') as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)
            return output_path
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            print(f"Retrying {url}... (Attempt {attempt + 1}/{max_retries})")

def merge_ts_files(ts_files, output_file):
    """
    Merges .ts files using ffmpeg.
    """
    try:
        with open("filelist.txt", 'w') as f:
            for ts_file in ts_files:
                f.write(f"file '{os.path.abspath(ts_file)}'\n")

        # Merge using ffmpeg
        subprocess.run(
            ['ffmpeg', '-f', 'concat', '-safe', '0', '-i', 'filelist.txt', '-c', 'copy', output_file],
            check=True
        )
    finally:
        if os.path.exists("filelist.txt"):
            os.remove("filelist.txt")

def download_and_merge_m3u8(m3u8_url, output_filename, output_folder=".", max_workers=10):
    """
    Downloads and merges M3U8 video segments.
    """
    # Ensure output folder exists
    os.makedirs(output_folder, exist_ok=True)

    if not output_filename.endswith(('.mp4', '.mkv', '.avi')):
        output_filename += ".mp4"

    print(f"Fetching M3U8 playlist from: {m3u8_url}")
    try:
        response = requests.get(m3u8_url, timeout=10)
        response.raise_for_status()
        playlist = m3u8.loads(response.text)
    except Exception as e:
        print(f"Failed to fetch M3U8 playlist: {e}")
        return

    base_uri = playlist.base_uri if playlist.base_uri else m3u8_url.rsplit('/', 1)[0] + '/'
    ts_urls = [urljoin(base_uri, segment.uri) for segment in playlist.segments]
    
    if not ts_urls:
        print("No segments found in the playlist.")
        return

    print(f"Found {len(ts_urls)} segments.")

    # Download TS segments with progress bar
    ts_files = []
    temp_dir = os.path.join(output_folder, "temp_ts_files")
    os.makedirs(temp_dir, exist_ok=True)

    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_url = {
                executor.submit(
                    download_ts_segment, 
                    url, 
                    os.path.join(temp_dir, f"{i:05d}.ts")
                ): url for i, url in enumerate(ts_urls)
            }
            
            for future in tqdm(concurrent.futures.as_completed(future_to_url), 
                              total=len(ts_urls), 
                              desc="Downloading TS segments"):
                try:
                    ts_file = future.result()
                    ts_files.append(ts_file)
                except Exception as e:
                    print(f"\nSegment download failed: {e}")

        # Sort files to ensure correct order
        ts_files.sort()

        # Merge TS files
        output_path = os.path.join(output_folder, output_filename)
        print(f"Merging TS files into {output_path}")
        merge_ts_files(ts_files, output_path)
        print(f"Video successfully saved to {output_path}")

    finally:
        # Clean up temporary files
        for ts_file in ts_files:
            try:
                os.remove(ts_file)
            except:
                pass
        try:
            os.rmdir(temp_dir)
        except:
            pass

if __name__ == "__main__":
    m3u8_url = input("Enter the M3U8 URL: ").strip()
    output_filename = input("Enter the output filename (e.g., video.mp4): ").strip()
    output_folder = input(f"Enter output folder (leave empty for current directory): ").strip() or "."
    
    download_and_merge_m3u8(m3u8_url, output_filename, output_folder)